Skip to content

Support Grain with Tfrecord in train_tokenizer#3522

Merged
copybara-service[bot] merged 1 commit intomainfrom
aireen/train_tokenizer_tfrecord
Apr 2, 2026
Merged

Support Grain with Tfrecord in train_tokenizer#3522
copybara-service[bot] merged 1 commit intomainfrom
aireen/train_tokenizer_tfrecord

Conversation

@aireenmei
Copy link
Copy Markdown
Collaborator

@aireenmei aireenmei commented Mar 30, 2026

Description

train_tokenizer.py already supports training with Grain in arrayrecord or parquet formats. Add support for tfrecord

Tests

Add TrainTokenizerFormatTest in tokenizer_test.py to test all 3 formats

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov
Copy link
Copy Markdown

codecov bot commented Mar 30, 2026

Codecov Report

❌ Patch coverage is 83.33333% with 2 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/trainers/tokenizer/train_tokenizer.py 83.33% 1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

@aireenmei aireenmei force-pushed the aireen/train_tokenizer_tfrecord branch from 94f2da2 to 30dd96e Compare April 2, 2026 00:18
@copybara-service copybara-service bot merged commit 0ee5a04 into main Apr 2, 2026
63 of 65 checks passed
@copybara-service copybara-service bot deleted the aireen/train_tokenizer_tfrecord branch April 2, 2026 21:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants